{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING W&B installed but not logged in.  Run `wandb login` or set the WANDB_API_KEY env variable.\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "\n",
    "import re\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "pd.set_option('max_rows', 100)\n",
    "\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()\n",
    "\n",
    "from simpletransformers.ner import NERModel, NERArgs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 文本清洗"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'raw_data/bmes_train.json' -> 'raw_data/bmes_train_corrected.json'\r\n"
     ]
    }
   ],
   "source": [
    "!cp -v 'raw_data/bmes_train.json' 'raw_data/bmes_train_corrected.json' \n",
    "\n",
    "# 这四个是多了中括号\n",
    "!sed -i 's/二河闸\\[/二河闸/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/克州\\[/克州/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/流域\\[/流域/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/徐水\\[-/徐水-/' 'raw_data/bmes_train_corrected.json'\n",
    "\n",
    "# 这些不规范 :(\n",
    "!sed -i 's/#河流\\*\\]//' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/岳西县#地区\\*，经\\[@/岳西县#地区*[@]，经/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/\\[@姜明/[@]姜明/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/义和村#地区\\*，全\\[@/义和村#地区*[@]，全[@]/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/，全\\[@\\]/，全/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/幸福乡苇子沟#地区\\*，经动力区\\[@黎明乡、/幸福乡苇子沟，经动力区黎明乡、/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/苇子沟#地区\\*，经动力区\\[@黎明乡-LOC/苇子沟-LOC\", \"黎明乡-LOC/' 'raw_data/bmes_train_corrected.json'\n",
    "!sed -i 's/遂城镇#地区\\*\\]/遂城镇/' 'raw_data/bmes_train_corrected.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_map = {\n",
    "    '河流': 'RIV',\n",
    "    '地区': 'LOC',\n",
    "    '水库': 'RES',\n",
    "    '水利术语': 'TER',\n",
    "}\n",
    "\n",
    "\n",
    "with open('raw_data/bmes_train_corrected.json') as f:\n",
    "    lines = f.read().split('\\n')\n",
    "    \n",
    "\n",
    "correct_lines = list()\n",
    "\n",
    "for line in lines:\n",
    "    if re.search('\"text\":', line):  # 文本\n",
    "        line = re.sub(r'#[^\\*]+\\*\\[@\\]', '', line)\n",
    "        \n",
    "    if re.search(r'^\\s+\"[^a-z]+.*\\*\\[@\\]', line):   # 标签\n",
    "        line = line.replace('*[@]', '\",\"')\n",
    "        line = line.replace('#', '-')\n",
    "        for label in label_map.keys():\n",
    "            line = line.replace(label, label_map[label]) \n",
    "    \n",
    "    correct_lines.append(line)\n",
    "    \n",
    "    \n",
    "with open('raw_data/bmes_train_corrected.json', 'w') as f:\n",
    "    f.write(\"\\n\".join(correct_lines))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 文本预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>text</th>\n",
       "      <th>entities</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>train_0</td>\n",
       "      <td>1976年冬，配合兴建涡河闸水利工程，开挖了涡河引河，全长1.2公里。</td>\n",
       "      <td>[涡河-RIV]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>train_1</td>\n",
       "      <td>宋代，甓社湖曾现珠光（河蚌珠光），并为在甓社湖居室临窗夜读的著名学者孙觉亲眼所见，被其好友沈...</td>\n",
       "      <td>[珠湖-LAK, 高邮湖-LAK, 甓社湖-LAK]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>train_2</td>\n",
       "      <td>2008年5月10日上午，山东省第五座跨黄河大桥，黄河首座———济南建邦黄河大桥在济南西外环...</td>\n",
       "      <td>[黄河-RIV, 济南-LOC, 山东省-LOC]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>train_3</td>\n",
       "      <td>学校曾获国家及省市级“职业教育先进单位”称号，是全国计算机应用技术考试培训基地（NIT），安...</td>\n",
       "      <td>[安徽省-LOC, 太湖县-LOC]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>train_4</td>\n",
       "      <td>东西两源汇合后，进入平原区，北流经过石埠嘴、船涨埠，至白洋淀后进入瓦埠湖。</td>\n",
       "      <td>[白洋淀-LAK, 瓦埠湖-LAK]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        id                                               text  \\\n",
       "0  train_0                1976年冬，配合兴建涡河闸水利工程，开挖了涡河引河，全长1.2公里。   \n",
       "1  train_1  宋代，甓社湖曾现珠光（河蚌珠光），并为在甓社湖居室临窗夜读的著名学者孙觉亲眼所见，被其好友沈...   \n",
       "2  train_2  2008年5月10日上午，山东省第五座跨黄河大桥，黄河首座———济南建邦黄河大桥在济南西外环...   \n",
       "3  train_3  学校曾获国家及省市级“职业教育先进单位”称号，是全国计算机应用技术考试培训基地（NIT），安...   \n",
       "4  train_4              东西两源汇合后，进入平原区，北流经过石埠嘴、船涨埠，至白洋淀后进入瓦埠湖。   \n",
       "\n",
       "                     entities  \n",
       "0                    [涡河-RIV]  \n",
       "1  [珠湖-LAK, 高邮湖-LAK, 甓社湖-LAK]  \n",
       "2   [黄河-RIV, 济南-LOC, 山东省-LOC]  \n",
       "3          [安徽省-LOC, 太湖县-LOC]  \n",
       "4          [白洋淀-LAK, 瓦埠湖-LAK]  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train = pd.read_json('raw_data/bmes_train_corrected.json')\n",
    "train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "4919it [00:01, 3132.84it/s]\n"
     ]
    }
   ],
   "source": [
    "train_data = list()\n",
    "\n",
    "for i, row in tqdm(train.iterrows()):\n",
    "    entities = row['entities']\n",
    "    text = row['text']\n",
    "    id_list = [i] * len(text)\n",
    "    label_list = [\"O\"] * len(text)\n",
    "    finished = list()\n",
    "    for entity in entities:\n",
    "        name, label = entity.split('-')\n",
    "        for m in re.finditer(name, text):     # 找到字符串的起始位置和它们所属的标签\n",
    "            start = m.start()\n",
    "            end = m.end()\n",
    "            if start not in finished:\n",
    "                label_list[start] = f'B_{label}'\n",
    "                finished.append(start)\n",
    "                for i in range(start+1, end):\n",
    "                    if i not in finished:\n",
    "                        label_list[i] = f'I_{label}'\n",
    "                        finished.append(i)\n",
    "    train_data.extend(zip(id_list, text, label_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sentence_id</th>\n",
       "      <th>words</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>6</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>年</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>0</td>\n",
       "      <td>冬</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0</td>\n",
       "      <td>，</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0</td>\n",
       "      <td>配</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0</td>\n",
       "      <td>合</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>0</td>\n",
       "      <td>兴</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   sentence_id words labels\n",
       "0            0     1      O\n",
       "1            0     9      O\n",
       "2            0     7      O\n",
       "3            0     6      O\n",
       "4            0     年      O\n",
       "5            0     冬      O\n",
       "6            0     ，      O\n",
       "7            0     配      O\n",
       "8            0     合      O\n",
       "9            0     兴      O"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_data = pd.DataFrame(\n",
    "    train_data, columns=[\"sentence_id\", \"words\", \"labels\"]\n",
    ")\n",
    "\n",
    "train_data.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_data[train_data.sentence_id==2354]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(263624, 3) 4919\n",
      "(237503, 3) (26121, 3)\n"
     ]
    }
   ],
   "source": [
    "print(train_data.shape, train_data.sentence_id.nunique())\n",
    "\n",
    "fold = 0\n",
    "eval_data = train_data[train_data['sentence_id'] % 10 == fold]\n",
    "train_data = train_data[train_data['sentence_id'] % 10 != fold]\n",
    "\n",
    "print(train_data.shape, eval_data.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args = NERArgs()\n",
    "model_args.train_batch_size = 8\n",
    "model_args.num_train_epochs = 8\n",
    "model_args.fp16 = False\n",
    "model_args.evaluate_during_training = True\n",
    "model_args.overwrite_output_dir = True\n",
    "model_args.save_steps = -1\n",
    "model_args.save_eval_checkpoints = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at hfl/chinese-roberta-wwm-ext were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at hfl/chinese-roberta-wwm-ext and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = NERModel(\"bert\", \n",
    "                 \"hfl/chinese-roberta-wwm-ext\",\n",
    "                 labels=train_data.labels.unique().tolist(),\n",
    "                 args=model_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e07850d521d480c98ba50b813ca9e6c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=4427.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ba724de25494e0f9ae19853759cc730",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Epoch', max=8.0, style=ProgressStyle(description_width='i…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0d1ac46763fb4433bc71739fecece7af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 0 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be45452f4c8a47469f72460d3a4fe264",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "685407c8c5e94f60bf409b67ba6cd182",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ff20e8c023b48aa9325a99e8a2cf3ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 1 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "952a5318f81f4e638524aa94f67be3dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ba3f6fabea94444b61e5bc8b000898a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7acfcb6f5790459d994baf9203ab4e39",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 2 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "388d5e9d150c4cf1a6c6855d159f51b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d2352a968c54437ab7a43a2089ade30d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c7bc929270347489e369be52037d4d1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 3 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df390708d5b8466b8fc4cb7022448988",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da6ea75412e6452aa89c32ff46356cd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6486b2b308d54df0861c6a515662370c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e032914fd8754ae8b3eeeec7f0b05116",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d28b2451ee1b4cd9b6c12dd500fa11b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 4 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0d343ef7eeb4ef6b9ea9f915ce5dc31",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eabe47faa485407982565143ce080700",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "986f250b550244a8a651e6e99c09bbd2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 5 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a554effdb60472fbeb5403765ee90e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b7695232352478e8a0e52ac5e528183",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5eb26956004448d38d9f873d71fc9570",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 6 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4455a9536e6244829eecb5cc83dfe0df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "12f93bd0020b4674bc23e59da6e2ae5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b23d2e0853f7417298ffc152c2ec67e3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Epoch 7 of 8', max=554.0, style=ProgressStyle(des…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfefdbe2a4f547079079a14840dab269",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "376b36610ec04867970feed1e02c5232",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "302bd072590746cc910d4cda517c4a65",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fd80ca34c2744bf892f09c50b39fb603",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(4432,\n",
       " {'global_step': [554, 1108, 1662, 2000, 2216, 2770, 3324, 3878, 4000, 4432],\n",
       "  'precision': [0.7627599243856332,\n",
       "   0.8141244421107903,\n",
       "   0.7480369515011548,\n",
       "   0.7753110514759698,\n",
       "   0.8034889687018985,\n",
       "   0.7756253065228053,\n",
       "   0.7747239263803681,\n",
       "   0.780535758171541,\n",
       "   0.779656862745098,\n",
       "   0.8037309511297951],\n",
       "  'recall': [0.8553259141494436,\n",
       "   0.8216746157922629,\n",
       "   0.8582405935347112,\n",
       "   0.8420773714891362,\n",
       "   0.8298887122416534,\n",
       "   0.838102808691044,\n",
       "   0.8365129835718071,\n",
       "   0.8415474297827239,\n",
       "   0.8428722840487546,\n",
       "   0.8105458399576047],\n",
       "  'f1_score': [0.8063952035973021,\n",
       "   0.8178821047078993,\n",
       "   0.7993583415597236,\n",
       "   0.8073161437825479,\n",
       "   0.8164754953076121,\n",
       "   0.8056546102903719,\n",
       "   0.8044336858198496,\n",
       "   0.80989417314803,\n",
       "   0.8100331041507511,\n",
       "   0.8071240105540897],\n",
       "  'train_loss': [0.058510277420282364,\n",
       "   0.2915749251842499,\n",
       "   0.07795960456132889,\n",
       "   0.033798232674598694,\n",
       "   0.01428133249282837,\n",
       "   0.4618372619152069,\n",
       "   0.032964661717414856,\n",
       "   0.0009222163935191929,\n",
       "   0.004802729468792677,\n",
       "   0.005723045207560062],\n",
       "  'eval_loss': [0.20857331072611193,\n",
       "   0.18421561739617778,\n",
       "   0.200805202126503,\n",
       "   0.22164959178095864,\n",
       "   0.23282309243034932,\n",
       "   0.25217998352262283,\n",
       "   0.28023308960180127,\n",
       "   0.3186245117216341,\n",
       "   0.3251513156198686,\n",
       "   0.33689696582094314]})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.train_model(train_data, eval_data=eval_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b72bfff583e84847abf24d43bf3f4155",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=492.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7c8b35c55cde4acdb5a4accf73b3d558",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Evaluation', max=62.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.33689696582094314,\n",
       " 'precision': 0.8037309511297951,\n",
       " 'recall': 0.8105458399576047,\n",
       " 'f1_score': 0.8071240105540897}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result, _, _ = model.eval_model(eval_data)\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型预测及生成提交文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>test_0</td>\n",
       "      <td>2013年有1个项目获得国家创新基金，2个项目获国家科技型中小企业创新基金支持，1家民营企业...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>test_1</td>\n",
       "      <td>汉江遥堤堤防防洪标准为防御汉江1964年型洪水，堤顶高程按设计洪水位加超高2米确定，面宽10...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>test_2</td>\n",
       "      <td>纳支流七里川河后称红岩河，主流经嘴头镇、白云乡、王家凌乡，于王家凌乡之擂鼓滩出本县境入留坝县。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>test_3</td>\n",
       "      <td>云山水库，位于七虎林河上游，1958年由复转官兵承建，水库的蓄水量为4750万立方米，控制了...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>test_4</td>\n",
       "      <td>武穴正处武山湖深水口古北江遗址上，原称青林湖，古长江九大要穴之一。</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       id                                               text\n",
       "0  test_0  2013年有1个项目获得国家创新基金，2个项目获国家科技型中小企业创新基金支持，1家民营企业...\n",
       "1  test_1  汉江遥堤堤防防洪标准为防御汉江1964年型洪水，堤顶高程按设计洪水位加超高2米确定，面宽10...\n",
       "2  test_2    纳支流七里川河后称红岩河，主流经嘴头镇、白云乡、王家凌乡，于王家凌乡之擂鼓滩出本县境入留坝县。\n",
       "3  test_3  云山水库，位于七虎林河上游，1958年由复转官兵承建，水库的蓄水量为4750万立方米，控制了...\n",
       "4  test_4                  武穴正处武山湖深水口古北江遗址上，原称青林湖，古长江九大要穴之一。"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test = pd.read_json('raw_data/bmes_test.json')\n",
    "test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "628it [00:00, 8990.67it/s]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "181c82c34a8b494ba5e17433f9f9e7ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=628.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77a1abbb75144b22995ff3a15a750d5f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Running Prediction', max=79.0, style=ProgressStyle(descri…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "test_data = list()\n",
    "\n",
    "for i, row in tqdm(test.iterrows()):\n",
    "    text = row['text']\n",
    "    test_data.append(text)\n",
    "    \n",
    "preds, raw_outputs = model.predict(test_data, split_on_space=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>entities</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>test_0</td>\n",
       "      <td>[20-LOC, 称号-TER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>test_1</td>\n",
       "      <td>[汉江-LOC, :3-TER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>test_2</td>\n",
       "      <td>[纳支-LOC, 坝县-TER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>test_3</td>\n",
       "      <td>[云山-LOC, 水量-TER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>test_4</td>\n",
       "      <td>[武穴-LOC, 之一-TER]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       id          entities\n",
       "0  test_0  [20-LOC, 称号-TER]\n",
       "1  test_1  [汉江-LOC, :3-TER]\n",
       "2  test_2  [纳支-LOC, 坝县-TER]\n",
       "3  test_3  [云山-LOC, 水量-TER]\n",
       "4  test_4  [武穴-LOC, 之一-TER]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub = pd.read_json('raw_data/bmes_sample.json')\n",
    "sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['七里川河-RIV', '红岩河-RIV', '嘴头镇-LOC', '白云乡-LOC', '王家凌乡-LOC', '留坝县-LOC']"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = preds[2]\n",
    "res_list = list() \n",
    "for char in pred:\n",
    "    if list(char.values())[0] != 'O':\n",
    "        res_list.extend(zip(char.keys(), char.values()))\n",
    "chars = [i[0] for i in res_list]\n",
    "labels = [i[1] for i in res_list]\n",
    "\n",
    "# 处理开头是 I 的情况\n",
    "offset = 0\n",
    "while offset < len(labels):\n",
    "    if labels[0].startswith('I_'):\n",
    "        labels.pop(0)\n",
    "        chars.pop(0)\n",
    "    else:\n",
    "        break\n",
    "    offset += 1\n",
    "\n",
    "words = list()\n",
    "word = ''\n",
    "for i, label in enumerate(labels):\n",
    "    if label.startswith('B_'):\n",
    "        if i > 0:\n",
    "            lbl = labels[i-1].replace('I_', '')\n",
    "            words.append(f'{word}-{lbl}')\n",
    "        word = chars[i]\n",
    "    else:\n",
    "        word += chars[i]\n",
    "        if i == len(labels) - 1:\n",
    "            lbl = labels[i].replace('I_', '')\n",
    "            words.append(f'{word}-{lbl}')\n",
    "ret = list()\n",
    "for item in words:\n",
    "    if item not in ret:\n",
    "        ret.append(item)\n",
    "                \n",
    "ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 628/628 [00:00<00:00, 22166.87it/s]\n"
     ]
    }
   ],
   "source": [
    "final = list()\n",
    "\n",
    "for pred in tqdm(preds):\n",
    "    res_list = list() \n",
    "    for char in pred:\n",
    "        if list(char.values())[0] != 'O':\n",
    "            res_list.extend(zip(char.keys(), char.values()))\n",
    "    chars = [i[0] for i in res_list]\n",
    "    labels = [i[1] for i in res_list]\n",
    "    # 处理开头是 I 的情况\n",
    "    offset = 0\n",
    "    while offset < len(labels):\n",
    "        if labels[0].startswith('I_'):\n",
    "            labels.pop(0)\n",
    "            chars.pop(0)\n",
    "        else:\n",
    "            break\n",
    "        offset += 1\n",
    "    words = list()\n",
    "    word = ''\n",
    "    for i, label in enumerate(labels):\n",
    "        if label.startswith('B_'):\n",
    "            if i > 0:\n",
    "                lbl = labels[i-1].replace('I_', '')\n",
    "                words.append(f'{word}-{lbl}')\n",
    "            word = chars[i]\n",
    "        else:\n",
    "            word += chars[i]\n",
    "            if i == len(labels) - 1:\n",
    "                lbl = labels[i].replace('I_', '')\n",
    "                words.append(f'{word}-{lbl}')\n",
    "    ret = list()\n",
    "    for item in words:\n",
    "        if item not in ret:\n",
    "            ret.append(item)\n",
    "    final.append(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>entities</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>test_0</td>\n",
       "      <td>[临河区-LOC, 巴彦淖尔市-LOC]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>test_1</td>\n",
       "      <td>[汉江-RIV]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>test_2</td>\n",
       "      <td>[七里川河-RIV, 红岩河-RIV, 嘴头镇-LOC, 白云乡-LOC, 王家凌乡-LOC...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>test_3</td>\n",
       "      <td>[云山水库-RES, 七虎林河-RIV, 上-B_TER, 上游-TER]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>test_4</td>\n",
       "      <td>[武山湖-LAK, 北江-RIV, 青林湖-LAK, 长江-RIV]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       id                                           entities\n",
       "0  test_0                               [临河区-LOC, 巴彦淖尔市-LOC]\n",
       "1  test_1                                           [汉江-RIV]\n",
       "2  test_2  [七里川河-RIV, 红岩河-RIV, 嘴头镇-LOC, 白云乡-LOC, 王家凌乡-LOC...\n",
       "3  test_3              [云山水库-RES, 七虎林河-RIV, 上-B_TER, 上游-TER]\n",
       "4  test_4                 [武山湖-LAK, 北江-RIV, 青林湖-LAK, 长江-RIV]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub['entities'] = final\n",
    "sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub.to_json('baseline.json', orient='records', force_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
